package cn.liu.knowledge.controller;

import lombok.AllArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingOptionsBuilder;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.InputStreamResource;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.List;

@RestController
@RequestMapping("document")
@AllArgsConstructor
public class DocumentController {

    private final VectorStore vectorStore;
    private final EmbeddingModel embeddingModel;

//    @SneakyThrows
    @PostMapping("doc")
    public void parseFile(@RequestParam MultipartFile file) throws IOException {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(new InputStreamResource(file.getInputStream()));
        TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
        List<Document> apply = tokenTextSplitter.apply(tikaDocumentReader.read());
        List<float[]> embeddings = this.embeddingModel.embed(apply, EmbeddingOptionsBuilder.builder().build(), new TokenCountBatchingStrategy());
        System.out.println(embeddings);
        vectorStore.add(apply);
    }
}
